{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "_execution_state": "idle",
    "_uuid": "051d70d956493feee0c6d64651c6a088724dca2a",
    "id": "VshpNbDMNXxq",
    "papermill": {
     "duration": 0.010774,
     "end_time": "2021-02-15T11:01:41.782833",
     "exception": false,
     "start_time": "2021-02-15T11:01:41.772059",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "# Penalized Linear Regressions: A Simulation Experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "UTUoszHwNXxr"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import random\n",
    "import numpy as np\n",
    "import warnings\n",
    "warnings.simplefilter('ignore')\n",
    "from sklearn.linear_model import LassoCV\n",
    "from sklearn.linear_model import RidgeCV, ElasticNetCV, LinearRegression, Lasso, Ridge\n",
    "from sklearn.metrics import r2_score\n",
    "from sklearn.model_selection import GridSearchCV\n",
    "import pandas as pd\n",
    "from sklearn.base import BaseEstimator\n",
    "random.seed(1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "boNconu7NXxs"
   },
   "source": [
    "## Data Generating Process"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "2-DeZXDfNXxs"
   },
   "source": [
    "We define a simple data generating process that allows for sparse, dense, and sparse+dense coefficients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "C16cAmFsAg_p"
   },
   "outputs": [],
   "source": [
    "def gen_data(n, p, *, regime=\"sparse\"):\n",
    "    # constants chosen to get R^2 of approximately .80\n",
    "    if regime == \"sparse\":\n",
    "        beta = ((1 / np.arange(1, p + 1)) ** 2) * 7\n",
    "    elif regime == \"dense\":\n",
    "        beta = ((np.random.normal(0, 1, p)) * 0.35)\n",
    "    elif regime == \"sparsedense\":\n",
    "        # taking out either results in an R^2 of approximately .69\n",
    "        beta = (((1 / np.arange(1, p + 1)) ** 2) * 5) + ((np.random.normal(0, 1, p)) * 0.25)\n",
    "\n",
    "    def true_fn(x):\n",
    "        return x @ beta\n",
    "\n",
    "    X = np.random.uniform(-.5, .5, size=(n, p))\n",
    "    gX = true_fn(X)\n",
    "    y = gX + np.random.normal(0, 1, size=n)\n",
    "    Xtest = np.random.uniform(-.5, .5, size=(n, p))\n",
    "    gXtest = true_fn(Xtest)\n",
    "    ytest = gXtest + np.random.normal(0, 1, size=n)\n",
    "    Xpop = np.random.uniform(-.5, .5, size=(100000, p))  # almost population limit\n",
    "    gXpop = true_fn(Xpop)\n",
    "    ypop = gXpop + np.random.normal(0, 1, size=100000)  # almost population limit\n",
    "    return X, y, gX, Xtest, ytest, gXtest, Xpop, ypop, gXpop, beta"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "kh_obVQKNXxt",
    "papermill": {
     "duration": 0.010616,
     "end_time": "2021-02-15T11:01:41.804126",
     "exception": false,
     "start_time": "2021-02-15T11:01:41.793510",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Data Generating Process: Approximately Sparse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cIV-ZeCGNXxt"
   },
   "outputs": [],
   "source": [
    "n = 100\n",
    "p = 400\n",
    "X, y, gX, Xtest, ytest, gXtest, Xpop, ypop, gXpop, betas = gen_data(n, p, regime=\"sparse\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "wOFX2dObNXxt"
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(r\"$Y$ vs. $g(X)$\")\n",
    "plt.scatter(gX, y)\n",
    "plt.xlabel(r\"$g(X)$\")\n",
    "plt.ylabel(r\"$Y$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "YdWNt-ZrNXxu"
   },
   "outputs": [],
   "source": [
    "print(f\"theoretical R^2:, {1 - np.var(ypop - gXpop) / np.var(ypop)}\")\n",
    "print(f\"theoretical R^2:, {np.var(gXpop) / np.var(ypop)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "lNBmiZkpJwSe"
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.scatter(range(len(betas)), abs(betas), s=5, color='b')\n",
    "plt.xlabel(r'$\\beta$')\n",
    "plt.ylabel('Magnitude (log scale)')\n",
    "plt.title(r'$\\beta$ Magnitude')\n",
    "plt.yscale('log')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "TEyFNuIoNXxv"
   },
   "source": [
    "## Lasso, Ridge, ElasticNet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "BuMLbuCaNXxv"
   },
   "source": [
    "We use sklearn's penalized estimators, which choose the penalty parameter via cross-validation (by default 5-fold cross-validation). These methods search over an adaptively chosen grid of hyperparameters. `ElasticNet` allows for a convex combination of `l1` and `l2` penalty and the ratio with `l1_ratio` corresponding to the proportion of the `l1` penalty."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "5t7kITTvNXxv"
   },
   "outputs": [],
   "source": [
    "# Regressions\n",
    "lcv = LassoCV().fit(X, y)\n",
    "ridge = RidgeCV().fit(X, y)\n",
    "enet = ElasticNetCV(l1_ratio=0.5).fit(X, y)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "-yFVOY9aNXxv"
   },
   "source": [
    "We calculate the R-squared on the small test set that we have"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "djcr2rq0NXxv"
   },
   "outputs": [],
   "source": [
    "r2_lcv = r2_score(ytest, lcv.predict(Xtest))\n",
    "r2_ridge = r2_score(ytest, ridge.predict(Xtest))\n",
    "r2_enet = r2_score(ytest, enet.predict(Xtest))\n",
    "r2_lcv, r2_ridge, r2_enet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "syiyl10ONXxv"
   },
   "source": [
    "We also calculate what the R-squared would be in the population limit (in our case for practical purposes when we have a very very large test sample)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "d6DZheRTNXxv"
   },
   "outputs": [],
   "source": [
    "r2_lcv = r2_score(ypop, lcv.predict(Xpop))\n",
    "r2_ridge = r2_score(ypop, ridge.predict(Xpop))\n",
    "r2_enet = r2_score(ypop, enet.predict(Xpop))\n",
    "r2_lcv, r2_ridge, r2_enet"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ioo1FK-kNXxw"
   },
   "source": [
    "We can also try this with fitting OLS after Lasso selects variables, but note, this is the wrong post-lasso OLS with cross-validation!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BEp9fWfHNXxw"
   },
   "outputs": [],
   "source": [
    "class PostLassoOLS:\n",
    "\n",
    "    def fit(self, X, y):\n",
    "        lasso = LassoCV().fit(X, y)\n",
    "        self.feats_ = np.abs(lasso.coef_) > 1e-6\n",
    "        self.lr_ = LinearRegression().fit(X[:, self.feats_], y)\n",
    "        return self\n",
    "\n",
    "    def predict(self, X):\n",
    "        return self.lr_.predict(X[:, self.feats_])\n",
    "\n",
    "    @property\n",
    "    def coef_(self):\n",
    "        return self.lr_.coef_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tcreqHGDNXxw"
   },
   "outputs": [],
   "source": [
    "plols = PostLassoOLS().fit(X, y)\n",
    "r2_score(ypop, plols.predict(Xpop))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "JZU3DVkgNXxw"
   },
   "source": [
    "## Plug-in Hyperparameter Lasso and Post-Lasso OLS"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "VdUmbg1bNXxx",
    "papermill": {
     "duration": 0.01429,
     "end_time": "2021-02-15T11:01:45.388902",
     "exception": false,
     "start_time": "2021-02-15T11:01:45.374612",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "Here we compute the lasso and ols post lasso using plug-in choices for penalty levels."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "xLXh-KPp4ivz"
   },
   "source": [
    "\\We use \"plug-in\" tuning with a theoretically valid choice of penalty $\\lambda = 2 \\cdot c \\hat{\\sigma} \\sqrt{n} \\Phi^{-1}(1-\\alpha/2p)$, where $c>1$ and $1-\\alpha$ is a confidence level, and $\\Phi^{-1}$ denotes the quantile function. Under homoskedasticity, this choice ensures that the Lasso predictor is well behaved, delivering good predictive performance under approximate sparsity. In practice, this formula will work well even in the absence of homoskedasticity, especially when the random variables $\\epsilon$ and $X$ in the regression equation decay quickly at the tails.\n",
    "\n",
    "In practice, many people choose to use cross-validation, which is perfectly fine for predictive tasks. However, when conducting inference, to make our analysis valid we will require cross-fitting in addition to cross-validation. As we have not yet discussed cross-fitting, we rely on this theoretically-driven penalty in order to allow for accurate inference in the upcoming notebooks."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "8qqapxWmNXxx"
   },
   "source": [
    "We pull an analogue of R's rlasso. Rlasso functionality: it is searching the right set of regressors. This function was made for the case of ***p*** regressors and ***n*** observations where ***p >>>> n***. It assumes that the error is i.i.d. The errors may be non-Gaussian or heteroscedastic.\\\n",
    "The post lasso function makes OLS with the selected ***T*** regressors.\n",
    "To select those parameters, they use $\\lambda$ as variable to penalize\\\n",
    "**Funny thing: the function rlasso was named like that because it is the \"rigorous\" Lasso.**\\\n",
    "We find a Python code that tries to replicate the main function of hdm r-package. It was made by [Max Huppertz](https://maxhuppertz.github.io/code/). His library is this [repository](https://github.com/maxhuppertz/hdmpy). If not using colab, download its repository and copy this folder to your site-packages folder. In my case it is located here ***C:\\Python\\Python38\\Lib\\site-packages*** . We need to install this package ***pip install multiprocess***."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "u7tAphVB8fVV"
   },
   "outputs": [],
   "source": [
    "!git clone https://github.com/maxhuppertz/hdmpy.git\n",
    "!pip install multiprocess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "OZz-6CRZNXxx"
   },
   "outputs": [],
   "source": [
    "# We wrap the package so that it has the familiar sklearn API\n",
    "import hdmpy\n",
    "\n",
    "\n",
    "class RLasso(BaseEstimator):\n",
    "\n",
    "    def __init__(self, *, post=True):\n",
    "        self.post = post\n",
    "\n",
    "    def fit(self, X, y):\n",
    "        self.rlasso_ = hdmpy.rlasso(X, y, post=self.post)\n",
    "        return self\n",
    "\n",
    "    def predict(self, X):\n",
    "        return X @ np.array(self.rlasso_.est['beta']).flatten() + np.array(self.rlasso_.est['intercept'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z_8rwavpNXxx"
   },
   "outputs": [],
   "source": [
    "rlasso = RLasso(post=False).fit(X, y)\n",
    "rlasso_post = RLasso(post=True).fit(X, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NZ6en14_NXxx"
   },
   "outputs": [],
   "source": [
    "r2_rlasso = r2_score(ytest, rlasso.predict(Xtest))\n",
    "r2_rlasso_post = r2_score(ytest, rlasso_post.predict(Xtest))\n",
    "r2_rlasso, r2_rlasso_post"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "dp4jXG8eNXxx"
   },
   "outputs": [],
   "source": [
    "r2_rlasso = r2_score(ypop, rlasso.predict(Xpop))\n",
    "r2_rlasso_post = r2_score(ypop, rlasso_post.predict(Xpop))\n",
    "r2_rlasso, r2_rlasso_post"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "RhBRFMSbNXxy"
   },
   "source": [
    "## LAVA: Dense + Sparse Coefficients"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "Ed5p108HNXxy"
   },
   "source": [
    "Now let's try the LAVA estimator"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NukBqoVfNXxy"
   },
   "outputs": [],
   "source": [
    "# We construct an sklearn API estimator that implements the LAVA method\n",
    "\n",
    "\n",
    "class Lava(BaseEstimator):\n",
    "\n",
    "    def __init__(self, *, alpha1=1, alpha2=1, iterations=5):\n",
    "        self.alpha1 = alpha1  # l1 penalty\n",
    "        self.alpha2 = alpha2\n",
    "        self.iterations = iterations\n",
    "\n",
    "    def fit(self, X, y):\n",
    "        ridge = Ridge(self.alpha2).fit(X, y)\n",
    "        lasso = Lasso(self.alpha1).fit(X, y - ridge.predict(X))\n",
    "\n",
    "        for _ in range(self.iterations - 1):\n",
    "            ridge = ridge.fit(X, y - lasso.predict(X))\n",
    "            lasso = lasso.fit(X, y - ridge.predict(X))\n",
    "\n",
    "        self.lasso_ = lasso\n",
    "        self.ridge_ = ridge\n",
    "        return self\n",
    "\n",
    "    def predict(self, X):\n",
    "        return self.lasso_.predict(X) + self.ridge_.predict(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "skdYbLrSNXxy"
   },
   "outputs": [],
   "source": [
    "lava = GridSearchCV(Lava(), {'alpha1': np.logspace(-4, 4, 20), 'alpha2': np.logspace(-4, 4, 20)},\n",
    "                    scoring='r2', n_jobs=-1)\n",
    "lava.fit(X, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "bOLdzf7tNXxz"
   },
   "outputs": [],
   "source": [
    "lava.best_estimator_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "0ZYyTCh3NXxz"
   },
   "outputs": [],
   "source": [
    "r2_lava = r2_score(ytest, lava.predict(Xtest))\n",
    "r2_lava"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Y5GuiuJMNXxz"
   },
   "outputs": [],
   "source": [
    "r2_lava = r2_score(ypop, lava.predict(Xpop))\n",
    "r2_lava"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "T9X2dSkxNXxz"
   },
   "source": [
    "## Summarizing Results"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uGp1L54WNXxz"
   },
   "outputs": [],
   "source": [
    "df = pd.DataFrame({'LassoCV': [r2_lcv],\n",
    "                   'RidgeCV': [r2_ridge],\n",
    "                   'ElasticNetCV': [r2_enet],\n",
    "                   'RLasso': [r2_rlasso],\n",
    "                   'RLassoOLS': [r2_rlasso_post],\n",
    "                   'Lava': [r2_lava]}).T\n",
    "df.columns = ['Population R-squared']\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "xB5_jtumNXx0"
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Different Models for Approximately Sparse Regime\")\n",
    "# 45 degree line\n",
    "plt.plot([np.min(gXtest), np.max(gXtest)], [np.min(gXtest), np.max(gXtest)], color='black', linestyle='--')\n",
    "\n",
    "# different models\n",
    "plt.scatter(gXtest, ridge.predict(Xtest), marker='^', c='brown', s=5, label='Ridge')\n",
    "plt.scatter(gXtest, enet.predict(Xtest), marker='v', c='yellow', s=5, label='ENet')\n",
    "plt.scatter(gXtest, rlasso.predict(Xtest), marker='D', c='red', s=5, label='RLasso')\n",
    "plt.scatter(gXtest, rlasso_post.predict(Xtest), marker='o', c='green', s=5, label='RLasso Post')\n",
    "plt.scatter(gXtest, lcv.predict(Xtest), marker='<', c='blue', s=5, label='LassoCV')\n",
    "plt.scatter(gXtest, lava.predict(Xtest), marker='>', c='magenta', s=5, label='Lava')\n",
    "plt.legend(loc='lower right')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "I_Vf3rmIIGbm"
   },
   "source": [
    "## Data Generating Process: Dense Coefficients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kfTohbAiIMKH"
   },
   "outputs": [],
   "source": [
    "n = 100\n",
    "p = 400\n",
    "X, y, gX, Xtest, ytest, gXtest, Xpop, ypop, gXpop, betas = gen_data(n, p, regime=\"dense\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "SQ5V0IAMIMgI"
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(r\"$Y$ vs. $g(X)$\")\n",
    "plt.scatter(gX, y)\n",
    "plt.xlabel(r\"$g(X)$\")\n",
    "plt.ylabel(r\"$Y$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "jwFhvQhAIdR1"
   },
   "outputs": [],
   "source": [
    "print(f\"theoretical R^2:, {1 - np.var(ypop - gXpop) / np.var(ypop)}\")\n",
    "print(f\"theoretical R^2:, {np.var(gXpop) / np.var(ypop)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1ksRgX46Lj7y"
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.scatter(range(len(betas)), abs(betas), s=5, color='b')\n",
    "plt.xlabel(r'$\\beta$')\n",
    "plt.ylabel('Magnitude (log scale)')\n",
    "plt.title(r'$\\beta$ Magnitude')\n",
    "plt.yscale('log')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FpqSFZmtIdX4"
   },
   "outputs": [],
   "source": [
    "# Regressions\n",
    "lcv = LassoCV().fit(X, y)\n",
    "ridge = RidgeCV(alphas=(1, 10, 25, 50, 100)).fit(X, y)\n",
    "enet = ElasticNetCV(l1_ratio=0.5).fit(X, y)\n",
    "rlasso = RLasso(post=False).fit(X, y)\n",
    "rlasso_post = RLasso(post=True).fit(X, y)\n",
    "lava = GridSearchCV(Lava(), {'alpha1': np.logspace(-4, 4, 20), 'alpha2': np.logspace(-4, 4, 20)},\n",
    "                    scoring='r2', n_jobs=-1).fit(X, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "j-3fkIjDIpOg"
   },
   "outputs": [],
   "source": [
    "r2_lcv = r2_score(ypop, lcv.predict(Xpop))\n",
    "r2_ridge = r2_score(ypop, ridge.predict(Xpop))\n",
    "r2_enet = r2_score(ypop, enet.predict(Xpop))\n",
    "r2_rlasso = r2_score(ypop, rlasso.predict(Xpop))\n",
    "r2_rlasso_post = r2_score(ypop, rlasso_post.predict(Xpop))\n",
    "r2_lava = r2_score(ypop, lava.predict(Xpop))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "7-fwE0LiIpTn"
   },
   "outputs": [],
   "source": [
    "df = pd.DataFrame({'LassoCV': [r2_lcv],\n",
    "                   'RidgeCV': [r2_ridge],\n",
    "                   'ElasticNetCV': [r2_enet],\n",
    "                   'RLasso': [r2_rlasso],\n",
    "                   'RLassoOLS': [r2_rlasso_post],\n",
    "                   'Lava': [r2_lava]}).T\n",
    "df.columns = ['Population R-squared']\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "_h3-sAHVIdel"
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Different Models for Dense Regime\")\n",
    "# 45 degree line\n",
    "plt.plot([np.min(gXtest), np.max(gXtest)], [np.min(gXtest), np.max(gXtest)], color='black', linestyle='--')\n",
    "\n",
    "# different models\n",
    "plt.scatter(gXtest, ridge.predict(Xtest), marker='^', c='brown', s=5, label='Ridge')\n",
    "plt.scatter(gXtest, enet.predict(Xtest), marker='v', c='yellow', s=5, label='ENet')\n",
    "plt.scatter(gXtest, rlasso.predict(Xtest), marker='D', c='red', s=5, label='RLasso')\n",
    "plt.scatter(gXtest, rlasso_post.predict(Xtest), marker='o', c='green', s=5, label='RLasso Post')\n",
    "plt.scatter(gXtest, lcv.predict(Xtest), marker='<', c='blue', s=5, label='LassoCV')\n",
    "plt.scatter(gXtest, lava.predict(Xtest), marker='>', c='magenta', s=5, label='Lava')\n",
    "plt.legend(loc='lower right')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "fRMnSTE0NXx0",
    "papermill": {
     "duration": 0.018842,
     "end_time": "2021-02-15T11:02:51.941852",
     "exception": false,
     "start_time": "2021-02-15T11:02:51.923010",
     "status": "completed"
    },
    "tags": []
   },
   "source": [
    "## Data Generating Process: Approximately Sparse + Small Dense Part"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "kn_fmyZ0NXx0"
   },
   "outputs": [],
   "source": [
    "n = 100\n",
    "p = 400\n",
    "X, y, gX, Xtest, ytest, gXtest, Xpop, ypop, gXpop, betas = gen_data(n, p, regime=\"sparsedense\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "uGUmB3Z_IeFL"
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(r\"$Y$ vs. $g(X)$\")\n",
    "plt.scatter(gX, y)\n",
    "plt.xlabel(r\"$g(X)$\")\n",
    "plt.ylabel(r\"$Y$\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "D86K6_U2NXx0"
   },
   "outputs": [],
   "source": [
    "print(f\"theoretical R^2:, {1 - np.var(ypop - gXpop) / np.var(ypop)}\")\n",
    "print(f\"theoretical R^2:, {np.var(gXpop) / np.var(ypop)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "cJeL2ubQLmwY"
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.scatter(range(len(betas)), abs(betas), s=5, color='b')\n",
    "plt.xlabel(r'$\\beta$')\n",
    "plt.ylabel('Magnitude (log scale)')\n",
    "plt.title(r'$\\beta$ Magnitude')\n",
    "plt.yscale('log')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BalMbmgsNXx0"
   },
   "outputs": [],
   "source": [
    "# Regressions\n",
    "lcv = LassoCV().fit(X, y)\n",
    "ridge = RidgeCV().fit(X, y)\n",
    "enet = ElasticNetCV(l1_ratio=0.5).fit(X, y)\n",
    "rlasso = RLasso(post=False).fit(X, y)\n",
    "rlasso_post = RLasso(post=True).fit(X, y)\n",
    "lava = GridSearchCV(Lava(), {'alpha1': np.logspace(-4, 4, 20), 'alpha2': np.logspace(-4, 4, 20)},\n",
    "                    scoring='r2', n_jobs=-1).fit(X, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "42uZ1gIqNXx1"
   },
   "outputs": [],
   "source": [
    "r2_lcv = r2_score(ypop, lcv.predict(Xpop))\n",
    "r2_ridge = r2_score(ypop, ridge.predict(Xpop))\n",
    "r2_enet = r2_score(ypop, enet.predict(Xpop))\n",
    "r2_rlasso = r2_score(ypop, rlasso.predict(Xpop))\n",
    "r2_rlasso_post = r2_score(ypop, rlasso_post.predict(Xpop))\n",
    "r2_lava = r2_score(ypop, lava.predict(Xpop))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2bZbPplkNXx1"
   },
   "outputs": [],
   "source": [
    "df = pd.DataFrame({'LassoCV': [r2_lcv],\n",
    "                   'RidgeCV': [r2_ridge],\n",
    "                   'ElasticNetCV': [r2_enet],\n",
    "                   'RLasso': [r2_rlasso],\n",
    "                   'RLassoOLS': [r2_rlasso_post],\n",
    "                   'Lava': [r2_lava]}).T\n",
    "df.columns = ['Population R-squared']\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "T0cV196UNXx1"
   },
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.title(\"Different Models for Approximately Sparse + Dense Regime\")\n",
    "# 45 degree line\n",
    "plt.plot([np.min(gXtest), np.max(gXtest)], [np.min(gXtest), np.max(gXtest)], color='black', linestyle='--')\n",
    "\n",
    "# different models\n",
    "plt.scatter(gXtest, ridge.predict(Xtest), marker='^', c='brown', s=5, label='Ridge')\n",
    "plt.scatter(gXtest, enet.predict(Xtest), marker='v', c='yellow', s=5, label='ENet')\n",
    "plt.scatter(gXtest, rlasso.predict(Xtest), marker='D', c='red', s=5, label='RLasso')\n",
    "plt.scatter(gXtest, rlasso_post.predict(Xtest), marker='o', c='green', s=5, label='RLasso Post')\n",
    "plt.scatter(gXtest, lcv.predict(Xtest), marker='<', c='blue', s=5, label='LassoCV')\n",
    "plt.scatter(gXtest, lava.predict(Xtest), marker='>', c='magenta', s=5, label='Lava')\n",
    "plt.legend(loc='lower right')\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "hide_input": false,
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
